from easycore.common.registry import Registry
from easycore.common.config import CfgNode
import os
import hashlib
import pickle


class DatasetLoader:
    registry = Registry("dataset")
    dataset_config_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'configs', 'dataset')
    dataset_dump_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'output', 'dataset')

    @classmethod
    def register(cls, dataset_type_name=None, obj=None):
        return cls.registry.register(dataset_type_name, obj)

    @classmethod
    def get(cls, dataset_name_or_config):
        # parse dataset name or config
        if isinstance(dataset_name_or_config, str):
            dataset_config_path = os.path.join(cls.dataset_config_dir, f'{dataset_name_or_config}.yaml')
            dataset_config = CfgNode.open(dataset_config_path)
        else:
            dataset_config = dataset_name_or_config

        # build dataset
        dataset_id = cls.hash_dataset(dataset_config)
        dataset_dump_path = os.path.join(cls.dataset_dump_dir, f'{dataset_id}.pkl')
        if os.path.exists(dataset_dump_path):
            with open(dataset_dump_path, 'rb') as f:
                dataset = pickle.load(f)
        else:
            dataset = cls.registry.get(dataset_config['type']).build(dataset_config)
            os.makedirs(os.path.dirname(dataset_dump_path), exist_ok=True)
            with open(dataset_dump_path, 'wb') as f:
                pickle.dump(dataset, f)

        return dataset

    @classmethod
    def hash_dataset(cls, dataset_config):
        m = hashlib.sha256()
        m.update(bytes(str(dataset_config), encoding='utf-8'))
        dataset_id = m.hexdigest()
        return dataset_id
